use std::{sync::Arc, time::Duration};

use wtask_base::{
    encrypt::WTASK_KEY,
    error::WResult,
    machine::{WSession, WSpawn},
    tcp::{stream_mpscchannel_exchange_nokey, stream_read_token, stream_write_token, tcp_listen, TCPClientType, WTcpStream, TOKEN_ALIVE},
    export::{
        log_debug, log_info,
        select, tokio_spawn, mpsc_channel, MpscSender, MpscReceiver, TokioMutex, tokio_sleep
    }
};


type ChannelReceiver = MpscReceiver<Vec<u8>>;
type ChannelSender = MpscSender<Vec<u8>>;
type PairData = (ChannelSender, ChannelReceiver);


async fn stream_process_proxy(
    stream: WTcpStream,
    server_all: Arc<TokioMutex<WSession<MpscSender<String>>>>,
    session_all: Arc<TokioMutex<WSession<PairData>>>,
) -> WResult<()> {
    let addr = stream.0.peer_addr()?;
    let addr_now = format!("{}:{}", addr.ip(), addr.port());
    log_info!("Server Accepted from: {addr_now}");
    let server_token = stream_read_token(stream.0.as_ref(), WTASK_KEY).await?;
    // log_debug!("Server get server_token: {} {:?}", addr_now, server_token);
    // log_debug!("Servers: {:?}", server_all.lock().await);
    let t = server_all.lock().await.get(&server_token).cloned();
    if let Some(server_sender) = t {
        // is session
        // log_debug!("Server send session id: {} {}", addr_now, server_token);
        stream_write_token(stream.1.as_ref(), &TCPClientType::Session.to_string(), &server_token).await?;
        // session
        let session_token = stream_read_token(stream.0.as_ref(), &server_token).await?;
        // log_debug!("Server start session: {} {} {}", addr_now, server_token, session_token);
        let t = session_all.lock().await.remove(&session_token);
        if let Some((tx, rx)) = t {
            // server session
            stream_mpscchannel_exchange_nokey(&stream, tx, rx, 1024).await?;
        } else {
            // client session
            let (tx_peer1, rx) = mpsc_channel::<Vec<u8>>(1024);
            let (tx, rx_peer2) = mpsc_channel::<Vec<u8>>(1024);
            session_all.lock().await.insert(&session_token, (tx_peer1, rx_peer2));
            // let start new session
            server_sender.send(session_token.clone()).await?;
            // client waiting server session
            stream_mpscchannel_exchange_nokey(&stream, tx, rx, 1024).await?;
        };
        Ok(())
    } else {
        // server master
        // log_debug!("Server Waiting for: {} {}", addr_now, server_token);
        let (tx_server, mut rx_server) = mpsc_channel::<String>(1024);
        // machine alive
        let tx_server_alive = tx_server.clone();
        let server_all_alive = Arc::clone(&server_all);
        let server_token_alive = server_token.clone();
        tokio_spawn(async move {
            // 心跳，为了正常关闭
            let time_gap = Duration::from_secs(5);
            loop {
                tokio_sleep(time_gap).await;
                if tx_server_alive.send(TOKEN_ALIVE.to_string()).await.is_err() {
                    break;
                }
            }
            // log_debug!("server stoped: {}", server_token_alive);
            server_all_alive.lock().await.remove(&server_token_alive);
        });
        // log sender
        server_all.lock().await.insert(&server_token, tx_server);
        // is master
        // log_debug!("Server sending type for: {} {}", addr_now, server_token);
        stream_write_token(stream.1.as_ref(), &TCPClientType::Master.to_string(), &server_token).await?;
        // loop
        // log_debug!("Servers: {:?}", server_all.lock().await);
        while let Some(session_token) = rx_server.recv().await {
            let t = stream_write_token(stream.1.as_ref(), &session_token, &server_token).await;
            if t.is_err() {
                break;
            }
        }
        // log_debug!("server stoped: {}", server_token);
        server_all.lock().await.remove(&server_token);
        Ok(())
    }
}



pub async fn tcp_proxy(port: u32) -> WResult<()> {
    let listener = tcp_listen(port, "Bridge")?;
    let server_all = Arc::new(TokioMutex::new(WSession::default()));
    let session_all = Arc::new(TokioMutex::new(WSession::default()));
    let mut wspawn = WSpawn::default();
    let (_t, mut kill_rx) = wspawn.catch_pause_signal(async move {});
    loop {
        select! {
            Ok((stream, _)) = listener.accept() => {
                log_debug!("Server get a connect ...");
                let server_all_now = Arc::clone(&server_all);
                let session_all_now = Arc::clone(&session_all);
                wspawn.spawn(async move {
                    let stream = stream.into_split();
                    let _ = stream_process_proxy(stream, server_all_now, session_all_now).await;
                });
            }
            _ = kill_rx.changed() => {
                log_info!("Shutting bridge ...");
                break;
            }
        }
    }
    Ok(())
}
